can we verify if we explicitly seet the backend to npu, this would also work?

When a mask of shape [batch, seq_len] or [batch, 1, 1, seq_len] is passed, the operator fails with an error similar to:
get unsupported atten_mask shape, the shape is [B, 1, 1, S] – while only shapes like [B, N, S, S], [B, 1, S, S], [1, 1, S, S], or [S, S] are accepted.

The _native_npu_attention function operates correctly as it leverages _maybe_modify_attn_mask_npu to reshape the attention mask from [batch_size, seq_len_k] to [batch_size, 1, seq_len_q, seq_len_k]. This reshaped format is compatible with the NPU backend.

Reference:
Ascend NPU fusion attention API:
https://www.hiascend.com/document/detail/zh/Pytorch/730/apiref/torchnpuCustomsapi/docs/context/torch_npu-npu_fusion_attention.md

When a mask of shape [batch, seq_len] or [batch, 1, 1, seq_len] is passed, the operator fails with an error

just want to make sure we're on the same page, could you share a code example that would produce this error on npu? specifically, I;d like to know if you are running the default attention backend, i.e. without wrapping your model call inside with attention_backend("_naive_npu")

Yes, this code fixes the issue with the "native" backend. After the fix, it runs correctly with the "_native_attention" backend. Here's an example:

python
import torch
import torch_npu
from diffusers import ErnieImagePipeline
from diffusers.utils import load_image

pipe = ErnieImagePipeline.from_pretrained("/model_dir/ERNIE-Image", torch_dtype=torch.bfloat16)
pipe = pipe.to("npu")
generator = torch.Generator(device="npu")

prompt = "A black and white Chinese rural dog"
images = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    num_inference_steps=50,
    guidance_scale=5.0,
    generator=generator,
    use_pe=True,
).images
images[0].save("ernie-image-output.png")

However, I've found that when using native_npu_attention as the backend, there are still some issues with mask handling. I've pushed an additional commit to the previous PR. PR link: #13451 .
Add the following code to the example to enable the _native_npu_attention backend: pipe.transformer.set_attention_backend("_native_npu")
Note: When processing masks, we need to perform expansion validation when 4D masks are passed in, and apply mask inversion to meet the NPU interface requirements.